# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/5 16:07 
@Author  : ZhangShenao 
@File    : 3.多查询结果融合策略.py 
@Desc    : 使用自定义的ReciprocalRankFusionRetriever,基于rrf算法,实现多查询结果融合策略

多查询结果融合策略,就是在Multi-Query的基础上,按照一定的算法对结果文档进行排序(Ranking)和融合(Fusion)
最终按照权重返回top_k个文档结果
排序融合通常采用RRF(Reciprocal Rank Fusion)倒数排序融合算法
https://plg.uwaterloo.ca/~gvcormac/cormacksigir09-rrf.pdf
rrf算法会兼顾文档的排名和出现频率

"""
import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_pinecone import PineconeVectorStore

from multi_query_rrf import ReciprocalRankFusionRetriever

dotenv.load_dotenv()

# 基于Pinecone,创建基础的VectorStoreRetriever
embeddings = ZhipuAIEmbeddings(model="embedding-3")
vector_store = PineconeVectorStore(
    index_name="rag",  # 指定索引名称
    embedding=embeddings,  # 指定Embedding模型
    namespace="llm-ops",  # 指定Namespace
    text_key="text",  # 指定文本属性key
)
base_retriever = vector_store.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={
        "k": 3,
        "score_threshold": 0.5
    }
)

# 创建MultiQueryRetriever多查询检索器
prompt = ChatPromptTemplate.from_template(
    "你是一个AI语言模型助手。你的任务是生成给定用户问题的3个不同版本，以从向量数据库中检索相关文档。"
    "通过提供用户问题的多个视角，你的目标是帮助用户克服基于距离的相似性搜索的一些限制。"
    "请用换行符分隔这些替代问题。"
    "原始问题：{question}"
)
rrf_retriever = ReciprocalRankFusionRetriever.from_llm(
    retriever=base_retriever,  # 指定底层进行检索的Prompt
    llm=ChatZhipuAI(model="glm-4-flash", temperature=0),  # 指定用于生成相似提问的LLM,可以将temperature设置为0,降低模型出现幻觉的概率
    prompt=prompt,  # MultiQueryRetriever采用的生成相似提问的Prompt默认是英文的，可以替换成中文的Prompt
)
rrf_retriever.top_k = 3

# 执行检索
# ReciprocalRankFusionRetriever会基于rrf算法对检索到的文档进行排名,返回最终的top_k个
docs = rrf_retriever.invoke("关于LLM应用配置的文档有哪些？")
print(len(docs))
print(docs[0])
